import torch
import torch.nn as nn
from functools import partial
from mmengine import Config
from mmdet.registry import MODELS
from mmengine.runner import Runner
from torchviz import make_dot
from mmdet.utils import register_all_modules
register_all_modules()

if __name__ == '__main__':
    # 加载配置文件
    config_file = '/home/lixiang/.config/labelapp/output/YYWahrFLTiiQskmIS9REg/faster-rcnn_r50-caffe_fpn_ms-2x_coco.py'
    cfg = Config.fromfile(config_file)
    # cfg.val_dataloader.batch_size = 1

    # 构建模型
    model = MODELS.build(cfg.model)
    if torch.cuda.is_available():
        model = model.cuda()
    # model = revert_sync_batchnorm(model)
    model.eval()

    # 构建数据加载器
    data_loader = Runner.build_dataloader(cfg.val_dataloader)
    data_iter = iter(data_loader)
    data_batch = next(data_iter)

    # 确定输入尺寸（直接从数据加载器中获取一个样本）
    data = model.data_preprocessor(data_batch)
    input_size = data['inputs'].shape
    print(input_size)

    _forward = model.forward
    model.forward = partial(_forward, data_samples=data['data_samples'])

    # 前向传播以获得计算图
    dummy_input = torch.randn(input_size)  
    if torch.cuda.is_available():
        dummy_input = dummy_input.cuda()
    output = model(dummy_input)
    print(output)    
    
    # 假定选择第一个张量进行可视化
    tensor_to_visualize = output[0][0]
    dot = make_dot(tensor_to_visualize, params=dict(model.named_parameters()))
    
    # 保存为文件
    dot.format = 'png'
    dot.render("mmdetection_model_graph")
